package com.kamikaze.docidset.impl;

import java.io.IOException;
import java.util.List;

import org.apache.lucene.search.DocIdSet;
import org.apache.lucene.search.DocIdSetIterator;

public class OrDocIdSetIterator extends DocIdSetIterator {

  private final class Item
  {
    public final DocIdSetIterator iter;
    public int doc;
    public Item(DocIdSetIterator iter)
    {
      this.iter = iter;
      this.doc = -1;
    }
  }
  private int _curDoc;
  private final Item[] _heap;
  private int _size;
 
  OrDocIdSetIterator(List<DocIdSet> sets) throws IOException
  {
    _curDoc = -1;
    _heap = new Item[sets.size()];
    _size = 0;
    for(DocIdSet set : sets)
    {
      _heap[_size++] = new Item(set.iterator());
    }
    if(_size == 0) _curDoc = DocIdSetIterator.NO_MORE_DOCS;
  }
  
  @Override
  public final int docID() {
    return _curDoc;
  }

  @Override
  public final int nextDoc() throws IOException
  {
    if(_curDoc == DocIdSetIterator.NO_MORE_DOCS) return DocIdSetIterator.NO_MORE_DOCS;

    Item top = _heap[0];
    while(true)
    {
      DocIdSetIterator topIter = top.iter;
      int docid;
      if((docid = topIter.nextDoc())!=DocIdSetIterator.NO_MORE_DOCS)
      {
        top.doc = docid;
        heapAdjust();
      }
      else
      {
        heapRemoveRoot();
        if(_size == 0) return (_curDoc = DocIdSetIterator.NO_MORE_DOCS);
      }
      top = _heap[0];
      int topDoc = top.doc;
      if(topDoc > _curDoc)
      {
        return (_curDoc = topDoc);
      }
    }
  }

  @Override
  public final int advance(int target) throws IOException
  {
    if(_curDoc == DocIdSetIterator.NO_MORE_DOCS) return DocIdSetIterator.NO_MORE_DOCS;

    if(target <= _curDoc) target = _curDoc + 1;
     
    Item top = _heap[0];
    while(true)
    {
      DocIdSetIterator topIter = top.iter;
      int docid;
      if((docid = topIter.advance(target))!=DocIdSetIterator.NO_MORE_DOCS)
      {
        top.doc = docid;
        heapAdjust();
      }
      else
      {
        heapRemoveRoot();
        if (_size == 0) return (_curDoc = DocIdSetIterator.NO_MORE_DOCS);
      }
      top = _heap[0];
      int topDoc = top.doc;
      if(topDoc >= target)
      {
        return (_curDoc = topDoc);
      }
    }
  }
  
// Organize subScorers into a min heap with scorers generating the earlest document on top.
  /*
  private final void heapify() {
      int size = _size;
      for (int i=(size>>1)-1; i>=0; i--)
          heapAdjust(i);
  }
  */
  /* The subtree of subScorers at root is a min heap except possibly for its root element.
   * Bubble the root down as required to make the subtree a heap.
   */
  private final void heapAdjust()
  {
    final Item[] heap = _heap;
    final Item top = heap[0];
    final int doc = top.doc;
    final int size = _size;
    int i = 0;
    
    while(true)
    {
      int lchild = (i<<1)+1;
      if(lchild >= size) break;
      
      Item left = heap[lchild];
      int ldoc = left.doc;
      
      int rchild = lchild+1;
      if(rchild < size){
        Item right = heap[rchild];
        int rdoc = right.doc;
        
        if(rdoc <= ldoc)
        {
          if(doc <= rdoc) break;
          
          heap[i] = right;
          i = rchild;
          continue;
        }
      }
      
      if(doc <= ldoc) break;
      
      heap[i] = left;
      i = lchild;
    }
    heap[i] = top;
  }

  // Remove the root Scorer from subScorers and re-establish it as a heap
  private final void heapRemoveRoot()
  {
    _size--;
    if (_size > 0)
    {
      Item tmp = _heap[0];
      _heap[0] = _heap[_size];
      _heap[_size] = tmp; // keep the finished iterator at the end for debugging
      heapAdjust();
    }
  }

}
